#!/usr/bin/env python
# -*- coding:utf-8 -*-
# author : liuyu
# date : 2018/8/6 0006

# !/usr/bin/env python

from collections import namedtuple

from ansible.executor.task_queue_manager import TaskQueueManager
from ansible.inventory.manager import InventoryManager
from ansible.parsing.dataloader import DataLoader
from ansible.playbook.play import Play
from ansible.plugins.callback import CallbackBase
from ansible.vars.manager import VariableManager
import ansible.constants as C

C.HOST_KEY_CHECKING = False


def default_taskinfos():
    TaskInfos = {
        "hostinfo": [
            {
                "host": "0.0.0.0",
                "port": 22,
                "user": "root",
                "password": "root"
            },
        ],
        "taskinfo": [
            {
                "module": "setup",
                "args": ""
            },
        ]
    }
    return TaskInfos


# Create a callback object so we can capture the output
class ResultsCollector(CallbackBase):

    def __init__(self, *args, **kwargs):
        super(ResultsCollector, self).__init__(*args, **kwargs)
        self.host_ok = {}
        self.host_unreachable = {}
        self.host_failed = {}

    def v2_runner_on_unreachable(self, result):
        self.host_unreachable[result._host.get_name()] = result

    def v2_runner_on_ok(self, result, *args, **kwargs):
        self.host_ok[result._host.get_name()] = result

    def v2_runner_on_failed(self, result, *args, **kwargs):
        self.host_failed[result._host.get_name()] = result


class AnsibleApi(object):

    def __init__(self, TaskInfos):
        self.taskinfo = TaskInfos.get('taskinfo', default_taskinfos()["taskinfo"])
        self.hostinfo = TaskInfos.get('hostinfo', default_taskinfos()["hostinfo"])
        self.resultinfo = []

    def run(self):
        host_list = []
        [host_list.append(i.get("host", '0.0.0.0')) for i in self.hostinfo]
        Options = namedtuple('Options', ['connection', 'module_path', 'forks', 'remote_user',
                                         'private_key_file', 'ssh_common_args', 'ssh_extra_args', 'sftp_extra_args',
                                         'scp_extra_args', 'become', 'become_method', 'become_user', 'verbosity',
                                         'check',
                                         'diff'])
        # required for
        # https://github.com/ansible/ansible/blob/devel/lib/ansible/inventory/manager.py#L204
        sources = ','.join(host_list)
        if len(host_list) == 1:
            sources += ','

        # initialize needed objects
        loader = DataLoader()
        options = Options(connection='smart', module_path=['/usr/share/ansible'], forks=100,
                          remote_user=None, private_key_file=None, ssh_common_args=None, ssh_extra_args=None,
                          sftp_extra_args=None, scp_extra_args=None, become=None, become_method=None,
                          become_user=None, verbosity=None, check=False, diff=False, )

        passwords = dict(sshpass=None, becomepass=None)

        # create inventory and pass to var manager
        inventory = InventoryManager(loader=loader, sources=sources)
        variable_manager = VariableManager(loader=loader, inventory=inventory)
        for host in self.hostinfo:
            inventory.add_host(host=host.get('host'), port=host.get('port'))
            hostname = inventory.get_host(hostname=host.get('host'))
            variable_manager.set_host_variable(host=hostname, varname='ansible_ssh_pass', value=host.get('password'))
            variable_manager.set_host_variable(host=hostname, varname='ansible_ssh_user', value=host.get('user'))
            variable_manager.set_host_variable(host=hostname, varname='ansible_ssh_port', value=host.get('port'))
            # print("设置全局管理变量 %s"%host.get('host'))
        # create play with tasks
        for task in self.taskinfo:
            # tasks.append(dict(action=dict(module=task.get("module"), args=task.get("args"))))
            play_source = dict(
                name="Ansible API Play",
                hosts=host_list,
                gather_facts='no',
                # tasks=[dict(action=dict(module='command', args=dict(cmd='/usr/bin/uptime')))]
                # tasks=[dict(action=dict(module='shell', args='/usr/sbin/ip a'))]
                tasks=[dict(action=dict(module=task.get("module"), args=task.get("args")), register='shell_out')
                       ,dict(action=dict(module='debug', args=dict(msg='{{shell_out}}')))
                       ]
                # dict(action=dict(module='setup', args='')),
                #        dict(action=dict(module='setup', args=''),register='shell_out'),
                #        dict(action=dict(module='debug', args=dict(msg='{{shell_out.stdout}}')))

            )
            play = Play().load(play_source, variable_manager=variable_manager, loader=loader)
            # actually run it
            tqm = None
            callback = ResultsCollector()
            try:
                tqm = TaskQueueManager(
                    inventory=inventory,
                    variable_manager=variable_manager,
                    loader=loader,
                    options=options,
                    passwords=passwords,
                )
                tqm._stdout_callback = callback
                result = tqm.run(play)
            finally:
                if tqm is not None:
                    tqm.cleanup()

            # print("UP ***********")
            # print(callback.host_ok.items())
            for host, result in callback.host_ok.items():
                try:
                    # print('{0} >>> {1}'.format(host, result._result))
                    self.resultinfo.append({host: {"message": result._result['msg']['changed'], "code": 0,
                                                   "results":{"message":result._result['msg']['stdout'],"code":result._result['msg']['rc'],'shell_run_time':result._result['msg']['delta']}}})
                except:
                    # print('{0} >>> {1}'.format(host, result._result))
                    self.resultinfo.append({host: {"message": result._result['msg']['changed'], "code": 0,"results":{"message":result._result['msg']['ansible_facts'],"code": 0}}})

            # print("FAILED *******")
            for host, result in callback.host_failed.items():
                # print('{0} >>> {1}'.format(host, result._result))
                self.resultinfo.append({host: {"message": result._result['msg'], "code": 1,"results":{"message":result._result['stderr'],"code":result._result['rc']}}})

            # print("DOWN *********")
            for host, result in callback.host_unreachable.items():
                # print('{0} >>> {1}'.format(host, result._result))
                self.resultinfo.append({host: {"message": result._result['msg'], "code": 1,"results":{}}})

        return self.resultinfo



def getsysinfo(hostinfo):
    TaskInfos={}
    TaskInfos["hostinfo"] = hostinfo
    # print(TaskInfos)
    ansibleapi = AnsibleApi(TaskInfos)
    result = ansibleapi.run()
    resultinfo = []
    # print(result)
    for res in result:
        for k in res:
            data = res[k]["results"]["message"]
            code = res[k]["code"]
            sysinfo = {}
            disksize = 0
            if code != 0:
                resultinfo.append({k: res[k]})
            else:
                for i in data["ansible_devices"]:
                    if i[0:2] in ("vd", "ss", "sd"):
                        disksize = disksize + int(data["ansible_devices"][i]["sectors"]) * int(
                            data["ansible_devices"][i]["sectorsize"]) / 1024 / 1024 / 1024
                sysinfo['disk'] = "{}".format(str(disksize) + str(" GB"))
                sysinfo['mem'] = round(data['ansible_memtotal_mb'] / 1024)
                sysinfo['cpu'] = "{} * {}".format(data['ansible_processor_count'] * data["ansible_processor_cores"],
                                                  data['ansible_processor'][2])
                sysinfo['system'] = data['ansible_distribution'] + " " + data['ansible_distribution_version']
                resultinfo.append({k: {"message": sysinfo, "code": 0}})
    return resultinfo


def execshell(hostinfo, shellcmd):
    TaskInfos = {}
    TaskInfos["taskinfo"]=[{}]
    TaskInfos["hostinfo"] = hostinfo
    TaskInfos["taskinfo"][0]["module"] = "shell"
    TaskInfos["taskinfo"][0]["args"] = shellcmd
    ansibleapi = AnsibleApi(TaskInfos)
    result = ansibleapi.run()
    return result

# data=getsysinfo([{"host": "10.1.41.200","port": 65534,"user": "root","password": "root"}])
# print(data)
# [{'10.1.41.220': {'disk': '30.0 GB', 'mem': 2, 'cpu': '4 * Intel(R) Core(TM) i7-7500U CPU @ 2.70GHz', 'system': 'CentOS 7.2.1511'}}]

# data=execshell([{"host": "10.1.41.200","port": 65534,"user": "root","password": "root"}], "uptime")
# print(data)
# [{'10.1.41.220': ' 01:14:30 up  2:24,  2 users,  load average: 0.01, 0.02, 0.05'}]
'''
[{'10.1.41.200': {'message': 'Failed to connect to the host via ssh: ssh: connect to host 10.1.41.200 port 6554: Connection refused\r\n', 'code': -1}}]
[{'10.1.41.200': {'message': ' 11:35:32 up  1:54,  4 users,  load average: 0.13, 0.07, 0.05', 'code': 0}}]
[{'10.1.41.200': {'message': 'Authentication failure.', 'code': -1}}]
[{'10.1.41.201': {'message': 'Failed to connect to the host via ssh: ssh: connect to host 10.1.41.201 port 65534: No route to host\r\n', 'code': -1}}]
[{'10.1.41.200': {'message': 'non-zero return code', 'code': 1}}]
'''
